Make Decoding Functions Graph-compatible (with XLA Support!) #271

abheesht17 · 2022-07-13T18:10:53Z

Resolves #241

Partially resolves #277

Will have to think a bit more about Beam Search.

mattdangerw

Thank you! Left a few comments

mattdangerw · 2022-07-20T23:00:21Z

keras_nlp/utils/text_generation_test.py

+
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+        self.assertEqual(


would just using mode.predict work? that would still hit all the compiled function paths, and allow you to avoid all this dummy metric stuff, which is hard to read

Could also be good to test the call on a batched dataset (where batch size not statically known), and on a single constant input, as you are doing here.

Ah, man. Stupid me. Should have used model.predict :P

mattdangerw · 2022-07-20T23:13:07Z

keras_nlp/utils/text_generation.py

+        ),
+        body=one_step,
+        loop_vars=[prompt],
+        shape_invariants=[tf.TensorShape(shape_invariants)],


Can we just pass tf.TensorShape([None, None]) as the shape invariant? Generally we should support a static batch size of None, tf data does this by default after calling .batch() for example. Might simplify the code a bit.

mattdangerw · 2022-07-20T23:17:28Z

keras_nlp/utils/text_generation_test.py

+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        model = TestModel()
+        model.compile(metrics=[dummy_metric])


If you add jit_compile=True does the test still pass?

If yes, we should test this with both jit_compile=True and False, using https://docs.pytest.org/en/6.2.x/parametrize.html

If no, we should either try to fix things with jit_compilation, or make sure we track that on a follow up issue.

@abheesht17 Let's try add a test case for jit_compile=True, and we can run it on GPU. We recently add GPU test support in this repo.

It's not working with jit_compile = True. Complete error logs: https://p.ip.fi/2TNt.

Looks like it won't work with shape_invariants.

mattdangerw · 2022-07-20T23:18:00Z

Also re-beam search, separate PR sounds good!

abheesht17

@mattdangerw, thanks for the review! Addressed all comments, save the jit_compile one.

abheesht17 · 2022-07-21T18:49:16Z

keras_nlp/utils/text_generation_test.py

+
+        expected_outputs = tf.tile([[3], [0]], [1, max_length - 2])
+        expected_outputs = tf.concat([inputs, expected_outputs], axis=1)
+        self.assertEqual(


Ah, man. Stupid me. Should have used model.predict :P

abheesht17 · 2022-07-22T14:32:47Z

keras_nlp/utils/text_generation_test.py

+
+        inputs = tf.constant([[0, 1], [1, 2]])
+        model = TestModel()
+        model.compile(metrics=[dummy_metric])


It's not working with jit_compile = True. Complete error logs: https://p.ip.fi/2TNt.

Looks like it won't work with shape_invariants.

mattdangerw · 2022-07-22T21:19:55Z

/gcbrun

mattdangerw · 2022-07-22T21:53:03Z

I think a pull request went by recently where we stopped doing seeded random generation because of discrepancies.

#269

Is this safe to land as is @chenmoneygithub @jessechancy ?

jessechancy · 2022-07-23T02:33:02Z

Seeded random generation should be removed. This is mainly because even when fully seeded, the randomness output is different on accelerator-testing with GPU.

mattdangerw

This is great. Just leaving some quick initial comments.

mattdangerw · 2022-08-10T17:25:28Z

keras_nlp/utils/text_generation.py

+            tf.cast(max_length, dtype=tf.int64),
+        ),
+        body=one_step,
+        loop_vars=[state],


can we avoid the state dict and just do loops_vars=(length, prompt) here? might be a little more readable

mattdangerw · 2022-08-10T18:24:54Z

keras_nlp/utils/text_generation.py

+
+    # Pad the prompt with `pad_token_id` to `max_length`. We use `map_fn` here
+    # because the batch_size might not be static.
+    prompt = tf.map_fn(


I feel like we should be able to make this simpler, we are just padding a batched tensor with pad_token_id to the sequence length right? We should not need a map_fn for this

mattdangerw · 2022-08-10T18:27:46Z

keras_nlp/utils/text_generation.py

+        loop_vars=[state],
+    )[0]
+
+    prompt = state["prompt"]
    if end_token_id is not None:
        prompt = mask_tokens_after_end_token(


do we even need this function anymore, if we are just starting with the correct sized tensor filled with pad_token_id?

hmm I guess we do to avoid random tokens after the end_token_id

mattdangerw

A few comments, but this looks pretty good to me! I only commented on one of the four utilities, but comments apply to all.

mattdangerw · 2022-08-11T17:40:59Z

keras_nlp/utils/text_generation.py

+    length = prompt.shape.as_list()[1]
+
+    # Pad the prompt with `pad_token_id` to `max_length`.
+    prompt = tf.concat(


nit: maybe split this into two lines for readability?

padding = tf.fill((tf.shape(prompt)[0], max_length - length), pad_token_id)
prompt = tf.concat((prompt, padding), axis=-1)

mattdangerw · 2022-08-11T17:45:17Z

keras_nlp/utils/text_generation.py

-    while i < max_length:
-        # If the prompt has reached our desired length, exit while loop.
-        pred = token_probability_fn(prompt)
+    length = prompt.shape.as_list()[1]


can we just do something like

batch_size, length = tf.shape(x)

And use that below? Then length and batch size are both tensors from the start.

Hmmm, I'll split this into two lines:

batch_size = tf.shape(prompt)[0] length = tf.shape(prompt)[1]

batch_size, length = tf.shape(x)

does not work in graph mode.

Stack trace: https://p.ip.fi/6YAg

ah destructuring is too fancy for autograph, I forgot. let's do

shape = tf.shape(prompt)
batch_size = shape[0]
length = shape[1]

mattdangerw · 2022-08-11T17:47:11Z

keras_nlp/utils/text_generation.py

+        return (length, prompt)
+
+    # Run a while loop till text of length `max_length` has been generated.
+    prompt = tf.while_loop(


length, prompt = tf.while_loop(...)

just to avoid that [1] which is not super readable

mattdangerw · 2022-08-11T17:48:54Z

keras_nlp/utils/text_generation_test.py

+
+        class TestModel(tf.keras.Model):
+            def call(self, inputs, training=False):
+                if not training:


is there a reason you have to do training switch here? it looks like you are never actually testing the training=True branch, might be nice to clean up the test a bit

mattdangerw · 2022-08-11T17:51:53Z

@chenmoneygithub do you know why the accelerator testing is failing here? This would be a great one to actually test on accelerators.

chenmoneygithub · 2022-08-11T17:59:04Z

I found it out, it's because the git branch has not synced to master branch, so the build file is outdated.

@abheesht17 Could you sync and push again? Thanks!

abheesht17 · 2022-08-11T18:00:10Z

I found it out, it's because the git branch has not synced to master branch, so the build file is outdated.

@abheesht17 Could you sync and push again? Thanks!

Sure!

mattdangerw

LGTM! Thanks!

chenmoneygithub

Nice work! Dropped a comment on the test.

Also could you help create a TODO(chenmoneygithub) at the top of text_generation.py saying we should refactor the code to reuse the same code? The padding + scatter_update handling is more complex than before, so it might be nice we can reuse the code.

chenmoneygithub · 2022-08-15T17:42:06Z

keras_nlp/utils/text_generation_test.py

@@ -342,7 +406,7 @@ def test_generate_with_ragged_prompt(self):
    def test_assert_probability_distribution_generation_is_correct(self):
        def token_probability_fn(inputs):
            batch_size = inputs.shape[0]
-            prob = tf.constant([[0.01, 0.01, 0.08, 0.9]])
+            prob = tf.constant([[0.0, 0.0, 0.0, 1.0]])


Do we need to change the number here? The original value seems to be more general?

Ah, yes. This was done to take care of accelerator testing. Seeded generation does not work, so, we've made the probability 1.

abheesht17 added 2 commits July 13, 2022 23:38

Make text gen functions graph compat

c2bb422

Minor edit

b2d258e

mattdangerw requested changes Jul 20, 2022

View reviewed changes

abheesht17 commented Jul 22, 2022

View reviewed changes

Address review comments - I

db54b85

Changes for XLA support

ea1cce3

abheesht17 changed the title ~~Make Decoding Functions Graph-compatible~~ Make Decoding Functions Graph-compatible (with XLA Support!) Aug 10, 2022

mattdangerw requested changes Aug 10, 2022

View reviewed changes

abheesht17 added 3 commits August 11, 2022 22:39

Address review comments - II

c58e21e

Some polishing up

ed5c944

Format code

49e198f

mattdangerw reviewed Aug 11, 2022

View reviewed changes

abheesht17 added 2 commits August 11, 2022 23:55

Address review comments - III

9c3caf6

Merge branch 'master' into text-gen-graph-ops

5e053a5

mattdangerw approved these changes Aug 11, 2022

View reviewed changes

Minor edit

da2f9a6

chenmoneygithub reviewed Aug 15, 2022

View reviewed changes

abheesht17 added 2 commits August 16, 2022 04:32

Add TODO comment

2bc4f25

Fix format

0cf4bee

chenmoneygithub approved these changes Aug 15, 2022

View reviewed changes

chenmoneygithub merged commit 34c0e27 into keras-team:master Aug 16, 2022

mattdangerw mentioned this pull request Sep 2, 2022

Use preallocated buffer for text generation #188

Closed

Make Decoding Functions Graph-compatible (with XLA Support!) #271

Make Decoding Functions Graph-compatible (with XLA Support!) #271

Uh oh!

Conversation

abheesht17 commented Jul 13, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mattdangerw commented Jul 20, 2022

Uh oh!

abheesht17 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mattdangerw commented Jul 22, 2022

Uh oh!

mattdangerw commented Jul 22, 2022

Uh oh!

jessechancy commented Jul 23, 2022

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mattdangerw commented Aug 11, 2022

Uh oh!

chenmoneygithub commented Aug 11, 2022

Uh oh!

abheesht17 commented Aug 11, 2022

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

chenmoneygithub left a comment

Choose a reason for hiding this comment

abheesht17 commented Jul 13, 2022 •

edited

Loading